import time
import os
import math
import argparse
from colorsys import rgb_to_hsv
from PIL import Image, ImageOps

IMG_INPUT_DIR = ''
IMG_OUTPUT_DIR = ''
SLICE_SIZE = 16


def zoom_func(img, zoom=1):
    width, height = img.size
    return ImageOps.fit(img, (width*zoom, height*zoom), Image.Resampling.LANCZOS)


def img_filter(suffix, img_name_list):
    print(f"开始排除后缀不符合的图片，原有{len(img_name_list)}张图片\n")
    count = 0
    for img_name in img_name_list:
        # print(img_name)
        if not img_name.endswith(suffix):
            count += 1
            print(f"找到{count}个后缀不符合要求的图片")
            img_name_list.remove(img_name)
    print(f"已从列表中去除非指定后缀的图片, 剩余{len(img_name_list)}张图片")
    return img_name_list


# 输入图片预处理：修改图片尺寸,并求图片的平均hsv值，保存至输出目录
def img_pretreatment(img_name_list):
    begin_time = time.time()
    count = 0.0
    for img_name in img_name_list:
        print(f"已预处理{math.floor(count / len(img_name_list) * 100)}%\n")
        # print(img_name)
        suffix = img_name.split('.')[-1]
        img = Image.open(IMG_INPUT_DIR + img_name)
        img = ImageOps.fit(img, (SLICE_SIZE, SLICE_SIZE), Image.LANCZOS)
        try:
            avg_h, avg_s, avg_v = get_avg_hsv(img)
        except IOError:
            print("跳过一张非RGB格式的图片\n")
            continue
        img.save(IMG_OUTPUT_DIR + str(avg_h) + '_' + str(avg_s) + '_' + str(avg_v) + '.' + suffix)
        count += 1
    time_cost = time.time() - begin_time
    print(f"全部预处理完毕，预处理耗时{time_cost}\n")


# 计算图片的平均hsv
def get_avg_hsv(img):
    h = 0
    s = 0
    v = 0
    width, height = img.size
    # img.load()函数不支持非RGB格式的图像，如png，会返回包含整数的 PixelAccess 对象
    pixels = img.load()
    count = 0
    for y in range(height):
        for x in range(width):
            pixel = pixels[x, y]
            if type(pixel) is int:
                raise IOError("图片非RGB格式，无法处理")
            r = pixel[0]
            g = pixel[1]
            b = pixel[2]
            hsv = rgb_to_hsv(r, g, b)
            h += hsv[0]
            s += hsv[1]
            v += hsv[2]
            count += 1
    h_avg = round(h / count, 3)
    s_avg = round(s / count, 3)
    v_avg = round(v / count, 3)
    return h_avg, s_avg, v_avg


# 从图片集中找hsv平均值最相似的一张图片
def find_closest(cur_sub_img, img_name_list):
    diff = 1000
    closest_img_name = None
    cur_avg_h, cur_avg_s, cur_avg_v = get_avg_hsv(cur_sub_img)
    for img_name in img_name_list:
        name = img_name.rsplit('.', 1)[0]
        avg_h = float(name.split('_')[0])
        avg_s = float(name.split('_')[1])
        avg_v = float(name.split('_')[2])
        now_diff = math.sqrt(math.pow(math.fabs(cur_avg_h - avg_h), 2) + math.pow(math.fabs(cur_avg_s - avg_s), 2) + math.pow(math.fabs(cur_avg_v - avg_v), 2))
        if now_diff < diff:
            diff = now_diff
            closest_img_name = img_name
    return closest_img_name


def generate_img(img, img_name_list):
    width, height = img.size
    print(f"原图 宽:{width}, 高:{height}")
    # 与img大小一样的背景图,用于贴图片块
    background = Image.new('RGB', img.size, (255, 255, 255))
    total_slice_num = math.floor(width * height / (SLICE_SIZE * SLICE_SIZE))
    print(f"共{total_slice_num}张子图")
    paste_complete = 0
    for y1 in range(0, height, SLICE_SIZE):  # 在img上每隔SLICE_SIZE取一个像素
        for x1 in range(0, width, SLICE_SIZE):
            x2 = x1 + SLICE_SIZE
            y2 = y1 + SLICE_SIZE
            # x1, y1为子图片位于img的锚点坐标(左上角), x2,y2为裁切终止坐标
            # 截取img的一个子图片用于从预处理后的图集中找平均hsv最匹配的图片
            slice_img = img.crop((x1, y1, x2, y2))
            closest_img_name = find_closest(slice_img, img_name_list)
            paste_img = Image.open(IMG_OUTPUT_DIR + closest_img_name)
            background.paste(paste_img, (x1, y1))
            paste_complete += 1
            print(f"已完成{int(float(paste_complete) / total_slice_num * 100)}%")
    print(f"处理了{paste_complete}张子图")
    return background


if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument("--input", type=str, required=False, default='img.jpg', help='input image')
    parse.add_argument("--zoom", type=int, required=False, default=1, help='image zoom')
    parse.add_argument("--db", type=str, required=False, default='imgs/', help="db relative path")
    parse.add_argument("--slice", type=int, required=False, default=32, help="slice size")
    parse.add_argument("--out", type=str, required=False, default='out/', help="output relative path")
    parse.add_argument("--skip", action='store_true', required=False, default=False, help="skip pretreatment")

    args = parse.parse_args()  # 解析参数
    image_name = args.input  # 输入的图片
    SLICE_SIZE = args.slice  # 每个子图的尺寸: SLICE_SIZE * SLICE_SIZE
    IMG_INPUT_DIR = args.db  # 如：'imgs/'
    IMG_OUTPUT_DIR = args.out  # 预处理后输出的路径，如：'out/'
    zoom = args.slice  # 对原始图像进行缩放的倍数
    skip = args.skip  # 输入的图片

    begin_time = time.time()  # 计时

    if not os.path.isdir(IMG_INPUT_DIR):
        print("找不到预处理的输入的目录")
        exit(-1)
    if not os.path.isdir(IMG_OUTPUT_DIR):
        if skip:
            print("找不到预处理输出目录，跳过预处理步骤无法进行后续的合成操作")
        else:
            os.mkdir(IMG_OUTPUT_DIR)  # 创建预处理输出目录

    img_name_list_input = None
    img_name_list_output = None
    if not skip:  # 不跳过预处理即进行预处理操作
        img_name_list_input = os.listdir(IMG_INPUT_DIR)

        img_name_list = img_filter('jpg', img_name_list_input)
        img_pretreatment(img_name_list_input)

    img_name_list_output = os.listdir(IMG_OUTPUT_DIR)

    rar_img = Image.open(image_name)
    rar_img = zoom_func(rar_img, zoom)
    out_img = generate_img(rar_img, img_name_list_output)
    img = Image.blend(out_img, rar_img, 0.5)  # 将out_img和out_img融合

    img.save('out1.jpg')
    out_img.save('out2.jpg')

    print("完毕")
    time_cost = time.time()-begin_time
    print(f"共耗时{int(time_cost)}秒，{int(time_cost/60)}分钟")
    out_img.show()

